import numpy as np
from PIL import Image
import torch
import torch.utils.data
import random
import os.path as osp
from pathlib import Path
from mynn.utils.img_util import rgb2ycbcr
from mynn.utils.registry import DATASET_REGISTRY
from mynn.utils import FileClient, imfrombytes, paired_random_crop, augment, img2tensor, mod_crop


@DATASET_REGISTRY.register()
class ImageDataset(torch.utils.data.Dataset):
    def __init__(self, dataset_opt, phase='train'):
        super().__init__()
        assert phase in ['train', 'test']

        self.dataset_opt = dataset_opt
        self.file_client = None
        self.phase = phase
        self.dataroot_gt = Path(dataset_opt['gt_root'])
        self.dataroot_lq = Path(dataset_opt['lq_root'])
        self.color_mode = dataset_opt['color_mode']
        self.mod_crop=dataset_opt['mod_crop']

        if self.phase == 'train':
            self.gt_size = dataset_opt['gt_size']
            self.scale = dataset_opt['scale']
            self.use_hflip = dataset_opt['use_hflip']
            self.use_rot = dataset_opt['use_rot']

        self.keys = []

        with open(dataset_opt['meta_info_file'], 'r') as fin:
            for line in fin:
                img_name, _ = line.split(' ')
                self.keys.append(f'{img_name}/')
        self.keys = sorted(self.keys)

    def get(self, index):
        return self.__getitem__(index)

    def __getitem__(self, index):
        if self.file_client is None:
            self.file_client = FileClient()

        img_name = self.keys[index]

        # Read GT image.
        img_gt_path = self.dataroot_gt / img_name
        img_bytes = self.file_client.get(img_gt_path)
        img_gt = imfrombytes(img_bytes, float32=True)

        if self.mod_crop > 0:
            img_gt = mod_crop(img=img_gt, scale=self.mod_crop)

        # Read LQ image.
        img_lq_path = self.dataroot_lq / img_name
        img_bytes = self.file_client.get(img_lq_path)
        img_lq = imfrombytes(img_bytes, float32=True)

        # Read bic image.
        img_bic = Image.open(img_lq_path)
        h, w = img_bic.height, img_bic.width
        img_bic = img_bic.resize((w * 4, h * 4), resample=Image.BICUBIC)
        img_bic = np.array(img_bic, dtype=np.float32) / 255.0

        if self.mod_crop > 0:
            img_lq = mod_crop(img=img_lq, scale=self.mod_crop)

        if self.color_mode == 'Y':
            img_gt = rgb2ycbcr(img_gt, y_only=True)
            img_lq = rgb2ycbcr(img_lq, y_only=True)

        if self.phase == 'train':
            # Randomly crop.
            img_gt, img_lq = paired_random_crop(img_gt, img_lq, self.gt_size,
                                                self.scale, img_gt_path)

            # Augmentation.
            img_results = augment([img_gt, img_lq], self.use_hflip,
                                  self.use_rot)
        elif self.phase == 'test':
            img_results = [img_gt, img_lq]

        img_results = img2tensor(img_results)
        img_gt = img_results[0]
        img_lq = img_results[1]

        return {'lq': img_lq, 'gt': img_gt, 'key': img_name, 'img_bic':img_bic}

    def __len__(self):
        return len(self.keys)
